import argparse
import torch
import logging
import sys
import os

# Add project root directory to path
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

from watermark.watermark import WatermarkEmbedder
import json

logging.basicConfig(format='[%(asctime)s] - %(message)s',
                    datefmt='%Y/%m/%d %H:%M:%S',
                    level=logging.INFO)

def main():
    parser = argparse.ArgumentParser(description='Verify Watermark in Transition Matrix')
    parser.add_argument('--matrix_path', type=str, required=True,
                       help='Path to transition matrix file')
    parser.add_argument('--watermark_key', type=str, required=True,
                       help='Watermark key to verify')
    parser.add_argument('--watermark_info_path', type=str, default=None,
                       help='Path to watermark info JSON file')
    parser.add_argument('--num_classes', type=int, default=None,
                       help='Number of classes (if watermark_info not provided)')
    parser.add_argument('--threshold', type=float, default=0.6,
                       help='Matching threshold for watermark verification (default: 0.6)')
    
    args = parser.parse_args()
    
    # Load transition matrix
    if not os.path.exists(args.matrix_path):
        logging.error(f"Transition matrix file does not exist: {args.matrix_path}")
        return False
    
    T = torch.load(args.matrix_path, weights_only=False)
    logging.info(f"Loaded transition matrix: {args.matrix_path}")
    logging.info(f"Transition matrix shape: {T.shape}")
    
    # Get number of classes
    if args.watermark_info_path and os.path.exists(args.watermark_info_path):
        with open(args.watermark_info_path, 'r') as f:
            watermark_info = json.load(f)
        num_classes = watermark_info['num_classes']
        expected_key = watermark_info.get('watermark_key', args.watermark_key)
        if expected_key != args.watermark_key:
            logging.warning(f"Warning: Provided watermark key does not match saved key")
            logging.warning(f"  Saved key: {expected_key}")
            logging.warning(f"  Provided key: {args.watermark_key}")
    else:
        num_classes = args.num_classes if args.num_classes else T.shape[0]
        logging.info(f"Using number of classes: {num_classes}")
    
    # Verify watermark
    embedder = WatermarkEmbedder(args.watermark_key)
    is_match, match_rate = embedder.verify_watermark(T, num_classes, args.threshold)
    
    # Output results
    logging.info("=" * 60)
    logging.info("Watermark Verification Results")
    logging.info("=" * 60)
    logging.info(f"Transition matrix path: {args.matrix_path}")
    logging.info(f"Watermark key: {args.watermark_key}")
    logging.info(f"Number of classes: {num_classes}")
    logging.info(f"Match rate: {match_rate:.4f} ({match_rate*100:.2f}%)")
    logging.info(f"Threshold: {args.threshold}")
    logging.info(f"Verification result: {'✓ Watermark detected' if is_match else '✗ Watermark not detected'}")
    logging.info("=" * 60)
    
    if is_match:
        logging.info("Conclusion: This model contains the expected watermark and can be traced to source.")
    else:
        logging.info("Conclusion: This model does not contain the expected watermark, or the watermark has been corrupted.")
    
    return is_match

if __name__ == '__main__':
    main()

